package com.peng.leetcode.tree;

import java.util.*;

/**
 * DelNodes
 *
 * @author: lupeng6
 * @create: 2021/2/23 15:14
 */
public class DelNodes {



    public static void main(String[] args) {

    }

    private Set<Integer> del = new HashSet<>();

    private List<TreeNode> answer = new ArrayList<>();

    public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        if (root == null) {
            return null;
        }
        for (int i : to_delete) {
            del.add(i);
        }
        if (!del.contains(root.val)) {
            answer.add(root);
        }
        dfs(root);
        return answer;
    }

    private TreeNode dfs(TreeNode node) {
        if (node == null) {
            return null;
        }
        node.left = dfs(node.left);
        node.right = dfs(node.right);
        if (del.contains(node.val)) {
            if (node.left != null) {
                answer.add(node.left);
            }
            if (node.right != null) {
                answer.add(node.right);
            }
            node = null;
        }
        return node;
    }
}
